from core.function import Function
import utils

# Max 类是一个用于计算数组中最大值的函数。
class Max(Function):
    def __init__(self, axis=None, keepdims=False):
        """
        初始化 Max 函数。

        参数:
            axis (None, int, tuple of int): 计算最大值的轴。默认为整个数组。
            keepdims (bool): 是否保持轴的维度。默认为 False。
        """
        self.axis = axis
        self.keepdims = keepdims

    def forward(self, x):
        """
        Max 函数的前向传播。

        参数:
            x (ndarray): 输入数组。

        返回:
            ndarray: 沿指定轴或整个数组的最大值。
        """
        y = x.max(axis=self.axis, keepdims=self.keepdims)
        return y

    def backward(self, gy):
        """
        Max 函数的反向传播。

        参数:
            gy (ndarray): 输出的梯度。

        返回:
            ndarray: 相对于输入的梯度。
        """

        x = self.inputs[0]
        y = self.outputs[0]()  # weakref

        shape = max_backward_shape(x, self.axis)
        gy = utils.functions_collect.reshape(gy, shape)
        y = utils.functions_collect.reshape(y, shape)
        cond = (x.data == y.data)
        gy = utils.functions_collect.broadcast_to(gy, cond.shape)
        return gy * cond

    def max(self, x):
        """
        计算最大值的接口。

        参数:
            x (ndarray): 输入数组。

        返回:
            ndarray: 沿指定轴或整个数组的最大值。
        """
        return self(x)


def max_backward_shape(x, axis):
    """
    计算 Max 函数反向传播的形状。

    参数:
        x (ndarray): 输入数组。
        axis (None, int, tuple of int): 计算最大值的轴。

    返回:
        list: 反向传播形状的列表。
    """
    if axis is None:
        axis = range(x.ndim)
    elif isinstance(axis, int):
        axis = (axis,)
    else:
        axis = axis

    shape = [s if ax not in axis else 1 for ax, s in enumerate(x.shape)]
    return shape

